from pathlib import Path
import json
from tqdm import tqdm
from copy import deepcopy

def merge_group(ori_path, save_path, group_nums=4):
    with open(ori_path, "r", encoding="utf8") as fp:
        ori_data = json.load(fp)
    new_data = deepcopy(ori_data)
    dataset_list = ori_data["dataset"]
    nums_dataset = len(dataset_list)
    print("===> nums_dataset: ", nums_dataset)
    new_dataset = []
    for i in tqdm(range(0, nums_dataset, group_nums)):
        new_resources = []
        for j in range(group_nums):
            try:
                new_resources.append(dataset_list[i+j]["collected_frames"][0]["resources"][0])
            except:
                pass
        new_dataset.append({
            "data_id": dataset_list[i]["data_id"],
            "collected_frames": [{"resources": new_resources}],
            "prelabeled_data": None
        })
    new_data["dataset"] = new_dataset

    with open(save_path, "w", encoding="utf8") as fp:
        json.dump(new_data, fp, indent=2)


if __name__ == "__main__":
    group_nums = 10
    ori_path = Path("w3_ap_parking_slot_ipm_5.0.0_coredata_1.json")
    save_path = f"{ori_path.stem}_group_{group_nums}{ori_path.suffix}"
    merge_group(ori_path, save_path, group_nums=group_nums)
